import os
import sys
from multiprocessing import Process, Queue



sys.path.append(os.path.abspath(".")) # 当前目录的绝对路径


def kwargs_to_cmd(kwargs):
    cmd = "python main_cv.py "
    for flag, val in kwargs.items():
        cmd += f"--{flag}={val} "

    return cmd


def run_exp(gpu_num, in_queue):
    while not in_queue.empty():
        try:
            experiment = in_queue.get(timeout = 3)
        except:
            return

        experiment["multigpu"] = gpu_num
        print(f"==> Starting experiment {kwargs_to_cmd(experiment)}")
        os.system(kwargs_to_cmd(experiment))


def main():

    dataset = 'splitcifar100'
    log_dir = f"./runs/{dataset}"

    methods = [
        # 'dense',
        # 'supsup',
        # 'testmask',
        # 'sswarmup',
        # 'ss'，
        'wm',
    ]

    experiments = []

    # # dense基线
    # # supsup基线
    # # sswarmup的入口

    # # 小论文 ss 3个随机种子
    # seeds = [0,1,2]
    # pruning_iters = [38, 36, 34, 29, 25, 22, 21, 18, 15, 13, 10, 8, 5, 3, 1]
    # for method in methods:
    #     config_fn = f'./experiments/config_{dataset}_{method}.yaml'
    #     for seed in seeds:
    #         for pruning_iter in pruning_iters:
    #             print('seed: ', seed, 'alpha: ', pruning_iter)
    #             kwargs = {
    #                 "config": config_fn,
    #                 "exp_name": f"{method}~seed={seed}",
    #                 "seed": seed,
    #                 "log-dir": log_dir,
    #                 "pruning_iter": pruning_iter,
    #                 "alpha": 0,
    #                 "lr_epochs": 170,
    #                 "lr1_min": 0.00001,
    #                 "train_score_lr": 0.001,
    #                 "train_weight_lr": 0.001
    #             }
    #             experiments.append(kwargs)

    # # 小论文 wm 3个随机种子
    # seeds = [0,1,2]
    # alphas = [0.00009,0.00007,0.00005,0.00003,0.00001,0.000009,0.000007,0.000005,0.000003,0.000001,0.0000007,0.0000001, 0]
    # for method in methods:
    #     config_fn = f'./experiments/config_{dataset}_{method}.yaml'
    #     for seed in seeds:
    #         for alpha in alphas:
    #             print('seed: ', seed, 'alpha: ', alpha)
    #             kwargs = {
    #                 "config": config_fn,
    #                 "exp_name": f"{method}~seed={seed}",
    #                 "seed": seed,
    #                 "log-dir": log_dir,
    #                 "alpha": alpha,
    #                 "lr_epochs": 170,
    #                 "lr1_min": 0.00001,
    #                 "train_score_lr": 0.001,
    #                 "train_weight_lr": 0.001
    #             }
    #             experiments.append(kwargs)

    queue = Queue()
    for e in experiments:
        queue.put(e)

    processes = []
    gpus = [0] # 只有一个gpu
    for gpu in gpus:
        p = Process(target=run_exp, args=(gpu, queue))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()


if __name__ == "__main__":
    main()